from typing import List
from torch.utils.data import Dataset
import pandas as pd
import numpy as np
import torch


class TurbineDataset(Dataset):
    '''
    base：表示时间序列预测中的“上一个点”
    diff：差分，特征
    y：目标值
    seq：历史窗口长度
    '''
    def __init__(self, base: List[pd.DataFrame], diff: List[pd.DataFrame], y: pd.Series, edge_index:torch.Tensor, edge_attr:torch.Tensor, seq:int, device):
        super(TurbineDataset, self).__init__()
        self.base = torch.stack([torch.tensor(base[b].to_numpy(dtype=np.float32), dtype=torch.float32).to(device) for b in range(len(base))])
        self.diff = torch.stack([torch.tensor(diff[d].to_numpy(dtype=np.float32), dtype=torch.float32).to(device) for d in range(len(diff))])
        self.y = torch.tensor(y.to_numpy(dtype=np.float32), dtype=torch.float32).to(device)
        self.edge_index = edge_index.to(device)
        self.edge_attr = edge_attr.unsqueeze(1).repeat(1, seq, 1).to(device)
        self.seq = seq
        self.device = device


    def __getitem__(self, index):
        return self.y[index + self.seq - 1], self.diff[:, index: index + self.seq, :], self.edge_index, self.edge_attr, self.y[index + self.seq]


    def __len__(self):
        return len(self.y) - self.seq